from dml import TripletModel, TripletDataset, X_train, y_train, triplet_loss
import torch
from torch.utils.data import DataLoader, Dataset
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import numpy as np

# 加载训练好的模型
model = TripletModel()
model.load_state_dict(torch.load('src/model_pth/triplet_model_iu_xray.pth',weights_only=False))
model.eval()

# 创建测试数据集
num_triplets = 100
test_dataset = TripletDataset(X_train, y_train, num_triplets)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 测试模型并收集嵌入向量
all_embeddings = []
all_labels = []
correct = 0
total = 0
with torch.no_grad():
    for batch in test_loader:
        anchor, positive, negative = batch
        
        # 确保输入数据的通道数为1
        anchor = anchor.unsqueeze(1)
        positive = positive.unsqueeze(1)
        negative = negative.unsqueeze(1)
        
        anchor_embedding = model(anchor)
        positive_embedding = model(positive)
        negative_embedding = model(negative)
        
        # 计算三元组损失
        loss = triplet_loss(anchor_embedding, positive_embedding, negative_embedding)
        
        # 计算准确率
        pos_dist = torch.sum(torch.square(anchor_embedding - positive_embedding), dim=1)
        neg_dist = torch.sum(torch.square(anchor_embedding - negative_embedding), dim=1)
        correct += torch.sum(pos_dist < neg_dist).item()
        total += anchor.size(0)
        
        # 收集嵌入向量和标签
        all_embeddings.append(anchor_embedding.numpy())
        all_labels.append(y_train[anchor.squeeze(1).squeeze(1).squeeze(1).long()])

# 计算准确率
accuracy = correct / total
print(f'Test Accuracy: {accuracy * 100:.2f}%')

# 将嵌入向量和标签转换为 numpy 数组
all_embeddings = np.vstack(all_embeddings)
all_labels = np.concatenate(all_labels)

# 使用 t-SNE 降维
tsne = TSNE(n_components=2, random_state=42)
embeddings_2d = tsne.fit_transform(all_embeddings)

# 可视化嵌入向量
plt.figure(figsize=(10, 10))
for label in np.unique(all_labels):
    indices = np.where(all_labels == label)
    plt.scatter(embeddings_2d[indices, 0], embeddings_2d[indices, 1], label=label)
plt.legend()
plt.show()